"""
AI数据分析 MCP
"""

from core.data_analyze.graph import build_graph
from fastmcp import FastMCP

graph = build_graph()

# 创建FastMCP实例
mcp = FastMCP("AIReport", host='0.0.0.0', port='8930')


@mcp.tool()
async def ai_report(task_description: str, datasets: list) -> str:
    """AI数据分析报告，根据用户的任务描述，自动查询数据并分析，生成数据分析报告。
    Args:
        task_description: 用户的分析任务描述。
        datasets: 任务相关的数据集
    Returns:
        str: 返回数据分析报告
    """
    response = await graph.ainvoke(
        input={
            "auto_decision": True,
            "max_retry_times": 3,
            "cur_retry_times": 0,
            "task_description": task_description,
            "datasets": datasets
        },
        stream_mode=["custom"])
    return response[-1][1]


if __name__ == "__main__":
    mcp.run(transport='sse')
